Add cross-platform Docker and hardware CI by haoruilee · Pull Request #2 · haoruilee/RL-Kernel

haoruilee · 2026-06-26T08:28:45Z

Adds CUDA and ROCm source-build images, label-gated hardware workflows, and local docs for running the same checks.\n\nThis fork PR is for validating the GitHub-hosted workflows before opening upstream.

gemini-code-assist

Code Review

This pull request introduces ROCm CI support, refactors the GPU CI scripts to be more configurable, updates Dockerfiles and project dependencies, and adds tests for CUDA SM90 backend prioritization. The reviewer feedback focuses on improving the robustness of the CI scripts, such as installing the correct optional dependencies (cuda and rocm) during testing, pre-installing flashinfer in the CUDA Docker image to avoid slow compilation, handling shell failures under set -euo pipefail when detecting Python, and preventing Git fetch failures when running scripts locally. Additionally, the reviewer suggests making ROCm test arguments configurable and robustifying the device capability mock in tests.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-26T08:31:14Z

+git fetch origin "$PR_SHA"
+git checkout --detach "$PR_SHA"
+"$PY" -m pip install -U pip setuptools wheel
+"$PY" -m pip install -e ".[test]"


The GPU CI script currently installs .[test], which omits the cuda optional dependencies (such as flashinfer and nvidia-ml-py). To ensure that the CUDA-specific optimized backends are installed and tested, we should install .[cuda,test] instead.

Suggested change

"$PY" -m pip install -e ".[test]"

"$PY" -m pip install -e ".[cuda,test]"

gemini-code-assist · 2026-06-26T08:31:14Z

+FLASH_AUTO_INSTALL="${RL_KERNEL_ROCM_FLASH_ATTN_AUTO_INSTALL:-1}"
+
+"$PY" -m pip install -U pip setuptools wheel
+"$PY" -m pip install -e ".[test]"


The ROCm CI script currently installs .[test], which omits the rocm optional dependencies (such as aiter). To ensure that the ROCm-specific optimized backends are installed and tested, we should install .[rocm,test] instead.

Suggested change

"$PY" -m pip install -e ".[test]"

"$PY" -m pip install -e ".[rocm,test]"

gemini-code-assist · 2026-06-26T08:31:14Z

+RUN python -m pip install -U pip setuptools wheel \
+ && python -m pip install -r requirements.txt \
+ && python -m pip install \
+    black \
+    isort \
+    mypy \
+    packaging \
+    pre-commit \
+    psutil \
+    pytest \
+    ruff


flashinfer does not publish pre-built wheels to PyPI directly, only source distributions. If it is not pre-installed in the Docker image, running pip install will attempt to compile it from source, which takes a very long time and often fails or times out in CI. Pre-installing it in the Docker image using the official wheel index for CUDA 12.4 and PyTorch 2.4 avoids this issue completely.

RUN python -m pip install -U pip setuptools wheel \ && python -m pip install -r requirements.txt \ && python -m pip install flashinfer -f https://flashinfer.ai/whl/cu124/torch2.4/index.html \ && python -m pip install \ black \ isort \ mypy \ packaging \ pre-commit \ psutil \ pytest \ ruff

gemini-code-assist · 2026-06-26T08:31:14Z

 PR_SHA="${PR_SHA:-$(date +%s)}"
-POD_NAME="rl-kernel-ci-${PR_SHA:0:7}"
-READY_RETRIES=60
+PROFILE_SLUG=$(printf "%s" "${RUNPOD_PROFILE_NAME:-gpu}" | tr -c "[:alnum:]-" "-")
+POD_NAME="rl-kernel-ci-${PR_SHA:0:7}-${PROFILE_SLUG}"


When running this script locally or manually without PR_SHA set, the fallback defaults to a Unix timestamp. This causes git fetch origin <timestamp> to fail later in the script because a timestamp is not a valid git ref. We can use a separate variable for the unique pod name while defaulting PR_SHA to main or HEAD for the git checkout.

Suggested change

PR_SHA="${PR_SHA:-$(date +%s)}"

POD_NAME="rl-kernel-ci-${PR_SHA:0:7}"

READY_RETRIES=60

PROFILE_SLUG=$(printf "%s" "${RUNPOD_PROFILE_NAME:-gpu}" | tr -c "[:alnum:]-" "-")

POD_NAME="rl-kernel-ci-${PR_SHA:0:7}-${PROFILE_SLUG}"

PR_SHA="${PR_SHA:-}"

if [ -z "$PR_SHA" ]; then

PR_SHA_FOR_POD=$(date +%s)

PR_SHA="main"

else

PR_SHA_FOR_POD="$PR_SHA"

fi

PROFILE_SLUG=$(printf "%s" "${RUNPOD_PROFILE_NAME:-gpu}" | tr -c "[:alnum:]-" "-")

POD_NAME="rl-kernel-ci-${PR_SHA_FOR_POD:0:7}-${PROFILE_SLUG}"

gemini-code-assist · 2026-06-26T08:31:15Z

 PY=$(command -v python3.11 || command -v python3)
 if [ -z "$PY" ]; then echo "[remote] FATAL: python not found in PATH"; exit 127; fi


Under set -euo pipefail, if neither python3.11 nor python3 is found, the command substitution $(command -v python3.11 || command -v python3) will return a non-zero exit code and cause the shell to exit immediately. This prevents the script from reaching the if [ -z "$PY" ] check and printing the helpful error message. Appending || true inside the command substitution resolves this.

Suggested change

PY=$(command -v python3.11 || command -v python3)

if [ -z "$PY" ]; then echo "[remote] FATAL: python not found in PATH"; exit 127; fi

PY=$(command -v python3.11 || command -v python3 || true)

if [ -z "$PY" ]; then echo "[remote] FATAL: python not found in PATH"; exit 127; fi

gemini-code-assist · 2026-06-26T08:31:15Z

+PY="${PYTHON:-$(command -v python3 || command -v python)}"
+if [ -z "$PY" ]; then
+  echo "[rocm-ci] FATAL: python not found in PATH"
+  exit 127
+fi


Similar to the GPU CI script, under set -euo pipefail, if neither python3 nor python is found, the command substitution will fail and cause the script to exit immediately, bypassing the helpful error message. Appending || true inside the command substitution ensures the error check is reached.

Suggested change

PY="${PYTHON:-$(command -v python3 || command -v python)}"

if [ -z "$PY" ]; then

echo "[rocm-ci] FATAL: python not found in PATH"

exit 127

fi

PY="${PYTHON:-$(command -v python3 || command -v python || true)}"

if [ -z "$PY" ]; then

echo "[rocm-ci] FATAL: python not found in PATH"

exit 127

fi

gemini-code-assist · 2026-06-26T08:31:15Z

+"$PY" -m pytest \
+  rl_engine/tests/test_dispatch.py \
+  tests/test_kernel_registry.py \
+  tests/test_attention_correctness.py \
+  tests/test_linear_logp.py \
+  tests/test_ratio_kl.py \
+  -q -rs


The list of test files is currently hardcoded in the ROCm CI script. Allowing PYTEST_ARGS to be overridden or appended to (similar to run_gpu_ci.sh) makes the script much more flexible for local development and custom CI runs.

Suggested change

"$PY" -m pytest \

rl_engine/tests/test_dispatch.py \

tests/test_kernel_registry.py \

tests/test_attention_correctness.py \

tests/test_linear_logp.py \

tests/test_ratio_kl.py \

-q -rs

PYTEST_ARGS="${PYTEST_ARGS:-rl_engine/tests/test_dispatch.py tests/test_kernel_registry.py tests/test_attention_correctness.py tests/test_linear_logp.py tests/test_ratio_kl.py -q -rs}"

"$PY" -m pytest $PYTEST_ARGS

gemini-code-assist · 2026-06-26T08:31:15Z

+    monkeypatch.delenv("RL_KERNEL_ROCM_ATTN_BACKEND", raising=False)
+    monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda")
+    monkeypatch.setattr(registry_module.device_ctx, "is_rocm", False)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability)


The mock for torch.cuda.get_device_capability is defined as lambda: capability, which takes 0 arguments. However, PyTorch's get_device_capability can accept a device argument (e.g., torch.cuda.get_device_capability(device)). Defining the mock with *args, **kwargs makes it more robust against future changes or other call sites.

Suggested change

monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability)

monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *args, **kwargs: capability)

Security (Critical): - rocm-self-hosted now checks out base branch (trusted scripts) instead of fork HEAD; run_rocm_ci.sh clones PR code into /tmp at runtime via PR_REPO_URL + PR_SHA — mirrors the cuda-runpod isolation pattern, so fork-controlled code never executes on the self-hosted runner directly. Correctness: - Restore torch.distributed.run for multi-GPU pytest: GPU_COUNT > 1 now runs `python -m torch.distributed.run --nproc_per_node=$GPU_COUNT -m pytest` so distributed code paths are actually exercised. - Install .[cuda,test,hf] instead of .[cuda,test] so transformers is present and test_stateless_hf_integration.py runs instead of silently skipping. Reliability: - Add set -e to run_gpu_ci.sh outer script (was set -uo pipefail only); pre-SSH failures now abort immediately instead of continuing with stale state. - Restore ENV TORCH_CUDA_ARCH_LIST="8.6" in Dockerfile.cuda so headless / CPU-only docker builds have a sensible default and don't fail or compile fat binaries for every arch. - Add cache-to to build-pr job in build-ci-image.yml so PR Docker builds warm the GHA cache; previously only main-branch pushes wrote the cache, leaving all PR builds cold. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Blocking: - Dockerfile.cuda: fix TORCH_CUDA_ARCH_LIST --build-arg passthrough (ENV ignored build-arg; now uses ARG+ENV pair) - pyproject.toml + setup.py: remove wrong 'aiter' PyPI package from rocm extra (that package is an unrelated async-iterator library); amd-aiter must be installed from source — see installation docs - ci/run_rocm_ci.sh: drop .[rocm,test] since rocm extra is now empty; install .[test] only - ci.yml: add CPU-only operator tests — test_linear_logp.py native tests and test_cpu_hal.py now run in every PR (regressions in NativeLinearLogpOp and HAL routing were previously undetected) Quality: - run_gpu_ci.sh: remove torch.distributed.run wrapping around pytest; tests use CUDA directly, not dist — running under distributed.run caused the full suite to execute once per rank (wasted cost and GPU contention) - run_gpu_ci.sh: tighten fallback POD_ID regex to stay anchored to "id" field, preventing accidental SHA/field matches - run_rocm_ci.sh: add informative error messages on git clone/fetch failure - Dockerfile.rocm: add ninja to pip installs (needed by flash-attn build) - docs/contributing/testing.md: fix Docker CUDA run example to use .[cuda,test]; add required secrets table (RUNPOD_API_KEY, RUNPOD_SSH_PRIVATE_KEY) with RunPod SSH key registration note - docs/getting_started/installation.md: add Testing section linking to contributing/testing.md so source-build contributors find the Docker docs Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

GitHub blocks GHA cache writes from fork-triggered pull_request events. The build-pr job only needs cache-from; cache-to belongs in build-and-push (trusted push-to-main context) where it already exists. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

haoruilee force-pushed the feat/cicd branch from f935530 to a6847ce Compare June 26, 2026 08:30

gemini-code-assist Bot reviewed Jun 26, 2026

View reviewed changes

haoruilee force-pushed the feat/cicd branch 2 times, most recently from ad05648 to 0b6b8f1 Compare June 26, 2026 09:30

ci: add cross-platform docker and hardware workflows

3d7f2ac

haoruilee force-pushed the feat/cicd branch from 0b6b8f1 to 3d7f2ac Compare June 26, 2026 09:44

haoruilee and others added 23 commits June 27, 2026 22:50

fix(ci): add strict gpu preflight safeguards

3ef3ba9

fix(ci): handle runpod create failures

2610c62

fix(ci): use runpod ssh key for pods

f3c5230

fix(ci): create runpod workspace before clone

5c6f18d

fix(ci): harden runpod pip installs

8924cd3

fix(ci): explicitly build cuda extension on runpod

a21a928

fix(ci): make flashinfer install optional for gpu core tests

7a03545

fix(ci): wait for runpod ssh readiness

daf021b

fix(cuda): correct sm90 tma barrier sequencing

337585e

fix(ci): retry transient runpod ssh disconnects

1524d2e

fix(cuda): pass sm90 tensor maps as raw addresses

55a2ac2

fix(cuda): use grid-constant addresses for sm90 tensor maps

43c6b0f

fix(cuda): store sm90 tensor maps in device memory

df29724

fix(cuda): use cuda12 tma tile modifier

e9ceeb8

fix(cuda): preserve shared address state for sm90 tma

53f30d6

fix(cuda): use cuda barriers for sm90 tma

3ebf596

fix(cuda): force global state for sm90 tensor maps

78a023a

fix(cuda): compile sm90 tma with cluster shared state

263e116

fix(cuda): fall back from sm90 logp for fp32 logits

734ea0f

fix(cuda): expose generic logp api from sm90 wrapper

d270e99

haoruilee closed this Jun 28, 2026

haoruilee pushed a commit that referenced this pull request Jul 7, 2026

docs: add WeChat community group #2 entry

bb05014

Navigation Menu

Repository navigation

Add cross-platform Docker and hardware CI - #2

Add cross-platform Docker and hardware CI#2
haoruilee wants to merge 24 commits into
mainfrom
feat/cicd

haoruilee commented Jun 26, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	"$PY" -m pip install -e ".[test]"
	"$PY" -m pip install -e ".[cuda,test]"

	"$PY" -m pip install -e ".[test]"
	"$PY" -m pip install -e ".[rocm,test]"

		PY=$(command -v python3.11 \|\| command -v python3)
		if [ -z "$PY" ]; then echo "[remote] FATAL: python not found in PATH"; exit 127; fi

	monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: capability)
	monkeypatch.setattr(torch.cuda, "get_device_capability", lambda args, *kwargs: capability)

Conversation

haoruilee commented Jun 26, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant